﻿#include <ntddk.h>

#include "driver.h"

#ifdef NDEBUG
#define LOG(...)
#else
#define LOG(...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, __VA_ARGS__)
#endif

NTSTATUS DriverCreateDispatch(
    IN struct _DEVICE_OBJECT* DeviceObject,
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);

    LOG("[Debug] DriverCreateDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverCreateDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverCloseDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject,
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);

    LOG("[Debug] DriverCloseDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverCloseDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverReadDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject, 
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);

    LOG("[Debug] DriverReadDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverReadDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverWriteDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject, 
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);

    LOG("[Debug] DriverWriteDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverWriteDispatch Leave\r\n");

    return status;
}

VOID 
DriverUnload(
    IN struct _DRIVER_OBJECT* DriverObject
    )
{
    PDEVICE_OBJECT deviceObject = NULL;
    PDEVICE_OBJECT nextDeviceObject = NULL;
    PDEVICE_EXTENSION deviceExtension = NULL;

    deviceObject = DriverObject->DeviceObject;
    while (deviceObject)
    {
        nextDeviceObject = deviceObject->NextDevice;
        deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;

        LOG("[Debug] deviceObject=%p, deviceExtension->DeviceObject=%p, deviceExtension->DeviceName=%wZ, deviceExtension->SymbolicLinkName=%wZ\r\n", deviceObject, deviceExtension->DeviceObject, &deviceExtension->DeviceName, &deviceExtension->SymbolicLinkName);

        IoDeleteSymbolicLink(&deviceExtension->SymbolicLinkName);
        IoDeleteDevice(deviceObject);
        deviceObject = nextDeviceObject;
    }
}

NTSTATUS
CreateDevice(
    IN PDRIVER_OBJECT DriverObject
    )
{
    NTSTATUS status;
    UNICODE_STRING deviceName;
    PDEVICE_OBJECT deviceObject;
    UNICODE_STRING symbolicLinkName;
    PDEVICE_EXTENSION deviceExtension;
    
    // 创建设备
    RtlInitUnicodeString(&deviceName, DEVICE_NAME);
    status = IoCreateDevice(
        DriverObject,
        sizeof(DEVICE_EXTENSION),
        &deviceName,
        FILE_DEVICE_UNKNOWN,
        0,
        FALSE,
        &deviceObject
    );
    if (!NT_SUCCESS(status))
    {
        LOG("[Error] IoCreateDevice deviceName=%wZ, status=%#X\r\n", &deviceName, status);
        return status;
    }
    deviceObject->Flags |= DO_BUFFERED_IO;

    // 创建符号链接
    RtlInitUnicodeString(&symbolicLinkName, DEVICE_SYMBOLIC_LINK_NAME);
    status = IoCreateSymbolicLink(&symbolicLinkName, &deviceName);
    if (!NT_SUCCESS(status))
    {
        IoDeleteDevice(deviceObject);
        LOG("[Error] IoCreateSymbolicLink symbolicLinkName=%wZ, deviceName=%wZ status=%#X\r\n", &symbolicLinkName, &deviceName, status);
        return status;
    }

    // 记录设备信息
    deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;
    deviceExtension->DeviceObject = deviceObject;
    deviceExtension->DeviceName = deviceName;
    deviceExtension->SymbolicLinkName = symbolicLinkName;
    
    return status;
}

NTSTATUS 
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
)
{
    LOG("[Debug] DriverObject DriverName=%wZ, HardwareDatabase=%wZ\r\n", &DriverObject->DriverName, DriverObject->HardwareDatabase);
    LOG("[Debug] RegistryPath=%wZ\r\n", RegistryPath);

    DriverObject->DriverUnload = DriverUnload;
    DriverObject->MajorFunction[IRP_MJ_CREATE] = DriverCreateDispatch;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverCloseDispatch;
    DriverObject->MajorFunction[IRP_MJ_READ] = DriverReadDispatch;
    DriverObject->MajorFunction[IRP_MJ_WRITE] = DriverWriteDispatch;

    CreateDevice(DriverObject);

    return STATUS_SUCCESS;
}